package etree

import (
	"regexp"
	"strings"
)

type (
	token struct {
		expr string
		tag  string
	}
)

var tokenizerRE, _ = regexp.Compile(`('[^']*'|"[^"]*"|::|//?|\.\.|\(\)|[/.*:\[\]\(\)@=])|((?:\{[^}]+\})?[^/\[\]\(\)@=\s]+)|\s+`)
var selectors = map[string]func([]Element, string, <-chan token) []Element{
	"":   selectChild,
	"*":  selectStar,
	".":  selectSelf,
	"//": selectDescendant,
	"[":  selectPredicate,
}

func tokenize(path string) <-chan token {
	tokens := make(chan token)
	go func() {
		for _, matches := range tokenizerRE.FindAllSubmatch([]byte(path), -1) {
			tokens <- token{expr: string(matches[1]), tag: string(matches[2])}
		}
		close(tokens)
	}()
	return tokens
}

func selectChild(es []Element, tag string, tokens <-chan token) []Element {
	result := make([]Element, 0)
	for _, e := range es {
		for i := 0; i < e.Length(); i++ {
			c := e.Item(i)
			if c.Tag() == tag {
				result = append(result, c)
			}
		}
	}
	return result
}

func selectStar(es []Element, tag string, tokens <-chan token) []Element {
	result := make([]Element, 0)
	for _, e := range es {
		for i := 0; i < e.Length(); i++ {
			c := e.Item(i)
			if c.Tag() != "" {
				result = append(result, c)
			}
		}
	}
	return result
}

func selectSelf(es []Element, tag string, tokens <-chan token) []Element {
	result := make([]Element, 0)
	for _, e := range es {
		result = append(result, e)
	}
	return result
}

func selectDescendant(es []Element, tag string, tokens <-chan token) []Element {
	token := <-tokens
	if token.expr == "*" {
		tag = "*"
	} else if token.tag != "" {
		tag = token.tag
	}
	result := make([]Element, 0)
	for _, e := range es {
		for d := range e.iter(tag) {
			if d != e {
				result = append(result, d)
			}
		}
	}
	return result
}

func selectPredicate(es []Element, tag string, tokens <-chan token) []Element {
	signature := ""
	key, value := "", ""
	result := make([]Element, 0)
	for {
		token := <-tokens
		if token.expr == "]" {
			break
		}
		if token.expr != "" && (strings.HasPrefix(token.expr[:1], `'`) || strings.HasPrefix(token.expr[:1], `"`)) {
			token.tag = token.expr[1 : len(token.expr)-1]
			token.expr = `'`
		}
		if token.expr == "" {
			signature += "-"
		} else {
			signature += token.expr
		}
		if token.tag != "" {
			if key == "" {
				key = token.tag
			} else {
				value = token.tag
			}
		}
	}
	switch signature {
	case "@-":
		for _, e := range es {
			if _, ok := e.Get(key); ok {
				result = append(result, e)
			}
		}
	case "@-='":
		for _, e := range es {
			if v, ok := e.Get(key); ok && v == value {
				result = append(result, e)
			}
		}
	}
	return result
}

func find(es []Element, path string) []Element {
	if path[len(path)-1:] == "/" {
		path = path + "*"
	}
	result := make([]Element, len(es))
	copy(result, es)
	tokens := tokenize(path)
	for token := range tokens {
		if selector, ok := selectors[token.expr]; ok {
			result = selector(result, token.tag, tokens)
		}
	}
	return result
}
